package com.cetc.sdp.kmga.cs.stream.nv;

import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.spark.streaming.api.java.JavaDStream;
import org.apache.spark.streaming.api.java.JavaInputDStream;
import org.apache.spark.streaming.api.java.JavaStreamingContext;
import org.apache.spark.streaming.kafka010.*;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
 * @description:
 * @author： DengQiang
 * @date: 2018/3/9 9:48
 */
public class KafkaDStream extends AbstractDStream {

    private List<String> topics;

    private Map<String, Object> kafkaParas;

    private volatile OffsetRange[] offsetRanges;

    private JavaInputDStream<ConsumerRecord<Object, Object>> kafkaStream;

    public KafkaDStream(String topic, Map<String, Object> kafkaParas) {
        this(Collections.singletonList(topic), kafkaParas);
    }

    public KafkaDStream(List<String> topics, Map<String, Object> kafkaParas) {
        if (topics == null || kafkaParas == null) {
            throw new NullPointerException("Kafka para can not be null");
        }
        this.topics = new ArrayList<>();
        this.topics.addAll(topics);
        this.kafkaParas = kafkaParas;
    }

    public JavaInputDStream<ConsumerRecord<Object, Object>> getKafkaDStream(JavaStreamingContext jssc) {
        if (kafkaStream == null) {
            kafkaStream = KafkaUtils.createDirectStream(
                    jssc,
                    LocationStrategies.PreferConsistent(),
                    ConsumerStrategies.Subscribe(topics, kafkaParas)
            );
        }
        //获取offsets
        kafkaStream.foreachRDD(javaRDD -> {
            offsetRanges = ((HasOffsetRanges) javaRDD.rdd()).offsetRanges();
        });
        return kafkaStream;
    }

    private void commitOffsets() {
        if (kafkaStream != null) {
            CanCommitOffsets offsets = (CanCommitOffsets)kafkaStream.inputDStream();
            offsets.commitAsync(offsetRanges);
//            kafkaStream.foreachRDD(javaRDD -> {
//                OffsetRange[] offsetRanges = ((HasOffsetRanges) javaRDD.rdd()).offsetRanges();
//                offsets.commitAsync(offsetRanges);
//            });
        }
    }



    @Override
    public JavaDStream<String> getDStreamAsPlaintext(JavaStreamingContext jssc) {
        return getKafkaDStream(jssc).map(record -> record.value().toString());
    }

    @Override
    public void afterBatchProcessed() {
        commitOffsets();
    }
}
